'''
给定多个dir
遍历其中的image

尝试 lotus 生成 depth
'''
image_types = ['.jpg','.jepg','.png','.bmp','.webp']
dirs = ['/mnt/nas/zhangshu/train_data/爬取_clean', # 2056
        '/mnt/nas/zhangshu/train_data/igpair/clothes', # 86873
        # '/mnt/nas/zhangshu/train_data/h-and-m-fashion-caption-12k/images', # 12437
        '/mnt/nas/zhangshu/train_data/dresscode/images', # 44263
        '/mnt/nas/zhangshu/train_data/viton/images', # 11647
        '/mnt/nas/zhangshu/train_data/clothes_collection_clean/images', # 16443
        ]
save_dir = '/mnt/nas/shengjie/datasets_depth2'
save_json_path = './depthredux_data.json'
import os
os.environ['CUDA_VISIBLE_DEVICES']='1'
import warnings
warnings.filterwarnings('ignore')
from os.path import splitext as ops
from os.path import join as osj
from os.path import exists as ose
from tqdm import tqdm
import shutil,pdb,json

if ose(save_dir):
    shutil.rmtree(save_dir)
os.makedirs(save_dir)

import torch

# from lotus.app_infer_depth import load_pipe_g,get_depth_by_lotus_g
# from lotus.app_infer_depth import load_pipe_d,get_depth_by_lotus_d

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# pipe = load_pipe_g()

def iter_all_images(dirs):
    total_count = 0
    for dir in dirs:
        assert os.path.exists(dir), dir

        # for root, _, _ in os.walk(dir):
        print(dir,total_count)
        with os.scandir(dir) as it:
            for entry in it:
                if not entry.is_file():
                    continue
                ext = os.path.splitext(entry.name)[-1].lower()
                if ext not in image_types:
                    continue
                
                total_count += 1
                if total_count % 10000 == 0:
                    print(dir,total_count)
                
                yield os.path.join(dir, entry.name)


save_data = {'data':[]}

# 使用 tqdm 包裹迭代器，而不是整个 filenames 列表
# count = 0
for img_path in iter_all_images(dirs):
    # count += 1
    # if count % 1000 == 0:
    #     print(count,end=' ',flush=True)
# for dir in dirs:
#     assert ose(dir),dir
#     for root,subdirs,filenames in os.walk(dir):
#         for filename in tqdm(filenames):
#             if not ops(filename)[-1] in image_types:continue
#             img_path = osj(root,filename)

            # depth = get_depth_by_lotus_g(pipe,img_path,None,device)

            # pdb.set_trace()

            # subdir = os.path.basename(root)
            # save_depth_dir = osj(save_dir,subdir)
            # os.makedirs(save_depth_dir,exist_ok=True)
            # save_depth_path = osj(save_depth_dir , ops(filename)[0]+'.png')
            
            # # pdb.set_trace()
            # depth.save(save_depth_path)

    save_data['data'].append({
        'img_path':img_path,
        # 'depth_path':save_depth_path
    })

with open(save_json_path,'w',encoding='utf-8') as f:
    json.dump( save_data , f , ensure_ascii=False , indent=4 )